import numpy as np
import numba as nb
from numba import cuda
from math import exp, ceil
from ..utils import get_grid_1d, get_grid_2d, get_grid_3d, assume_ndarray


@cuda.jit()
def _gpu_matmul_2d(a: np.ndarray, b: np.ndarray, c: np.ndarray):
    """二维矩阵乘法"""
    x, y = cuda.grid(2)
    if x < c.shape[0] and y < c.shape[1]:
        sum = 0
        for i in range(a.shape[1]):
            sum += a[x, i] * b[i, y]
        c[x, y] = sum


@cuda.jit()
def _gpu_matmul_3d(a: np.ndarray, b: np.ndarray, c: np.ndarray):
    """三维矩阵乘法"""
    x, t = cuda.grid(2)
    y = t // c.shape[2]
    z = t % c.shape[2]
    if x < c.shape[0] and t < c.shape[1] * c.shape[2]:
        sum = 0
        for i in range(a.shape[2]):
            sum += a[x, y, i] * b[x, i, z]
        c[x, y, z] = sum


@cuda.jit()
def _gpu_matmul_1d_3d(a: np.ndarray, b: np.ndarray, c: np.ndarray):
    x, y = cuda.grid(2)
    if x < c.shape[0] and y < c.shape[1]:
        sum = 0
        for i in range(b.shape[1]):
            sum += a[i] * b[x, i, y]
        c[x, y] = sum


@cuda.jit()
def _gpu_matmul_2d_1d(a: np.ndarray, b: np.ndarray, c: np.ndarray):
    x = cuda.grid(1)
    if x < c.shape[0]:
        sum = 0
        for i in range(a.shape[1]):
            sum += a[x, i] * b[i]
        c[x] = sum


def matmul(a: np.ndarray, b: np.ndarray):
    a, b = assume_ndarray(a, b)
    if a.ndim == 2 and b.ndim == 2:
        assert a.shape[1] == b.shape[0]
        c = np.empty((a.shape[0], b.shape[1]))

        grid, block = get_grid_2d(c.shape[0], c.shape[1])
        _gpu_matmul_2d[grid, block](a, b, c)

        return c
    elif a.ndim == 1:
        assert a.shape[0] == b.shape[-2]
        shape = b.shape[:-2] + (b.shape[-1],)

        b = b.reshape(-1, b.shape[-2], b.shape[-1])
        c = np.empty(shape).reshape(-1, b.shape[-1])

        grid, block = get_grid_2d(c.shape[0], c.shape[1])
        _gpu_matmul_1d_3d[grid, block](a, b, c)

        return c.reshape(shape)
    elif b.ndim == 1:
        assert a.shape[-1] == b.shape[0]
        shape = a.shape[:-1]

        a = a.reshape(-1, b.shape[0])
        c = np.empty((a.shape[0],))

        grid, block = get_grid_1d(c.shape[0])
        _gpu_matmul_2d_1d[grid, block](a, b, c)

        return c.reshape(shape)
    elif a.ndim > 2 or b.ndim > 2:
        assert a.shape[-1] == b.shape[-2]
        assert a.shape[:-2] == b.shape[:-2]
        shape = a.shape[:-1] + (b.shape[-1],)

        a = a.reshape(-1, a.shape[-2], a.shape[-1])
        b = b.reshape(-1, b.shape[-2], b.shape[-1])
        c = np.empty(shape).reshape(-1, a.shape[-2], b.shape[-1])

        grid, block = get_grid_2d(c.shape[0], c.shape[1] * c.shape[2])
        _gpu_matmul_3d[grid, block](a, b, c)

        return c.reshape(shape)
    else:
        raise Exception("参数错误")
